package com.hitqz.robot.biz.config;

import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

@Component
public class SseEmitterManager {

    /**
     * topic -> userId -> SseEmitter
     */
    private static final Map<String, ConcurrentMap<Long, CopyOnWriteArrayList<SseEmitter>>> emitters = new ConcurrentHashMap<>();



    /**
     * 添加SseEmitter
     */
    public SseEmitter addEmitter(String topic,Long userId){
        SseEmitter emitter =  new SseEmitter(60 * 60 * 1000L);

        emitter.onCompletion(() -> removeEmitter(topic, userId, emitter));
        emitter.onTimeout(() -> removeEmitter(topic, userId, emitter));
        emitter.onError((ex) -> removeEmitter(topic, userId, emitter));

        emitters.computeIfAbsent(topic, k -> new ConcurrentHashMap<>())
                .computeIfAbsent(userId, k -> new CopyOnWriteArrayList<>())
                .add(emitter);

        return emitter;

    }


    public void removeEmitter(String topic,Long userId, SseEmitter emitter) {
        Map<Long, CopyOnWriteArrayList<SseEmitter>> userMap = emitters.get(topic);
        if (userMap != null) {
            List<SseEmitter> list = userMap.get(userId);
            if (list != null) {
                list.remove(emitter);
                if (list.isEmpty()) {
                    userMap.remove(userId);
                }
            }
            if (userMap.isEmpty()) {
                emitters.remove(topic);
            }
        }
    }


    /**
     * 发送根据用户数据权限筛选后的数据
     * @param topic
     * @param userId
     * @param data
     */
    public void sendMessageByUser(String topic, Long userId, Object data) {
        Map<Long, CopyOnWriteArrayList<SseEmitter>> userMap = emitters.get(topic);
        if (userMap == null) return;

        List<SseEmitter> emitters = userMap.get(userId);
        if (emitters != null) {
            for (SseEmitter emitter : emitters) {
                try {
                    emitter.send(SseEmitter.event().data(data));
                } catch (IOException e) {
                    removeEmitter(topic, userId, emitter);
                }
            }
        }
    }


    public void sendMessage(String channel, Object data) {
        Map<Long, CopyOnWriteArrayList<SseEmitter>> userMap = emitters.get(channel);
        if (userMap == null) return;

        for (Map.Entry<Long, CopyOnWriteArrayList<SseEmitter>> entry : userMap.entrySet()) {
            List<SseEmitter> emitters = entry.getValue();
            for (SseEmitter emitter : emitters) {
                try {
                    emitter.send(SseEmitter.event().data(data));
                } catch (IOException e) {
                    removeEmitter(channel, entry.getKey(), emitter);
                }
            }
        }

    }


}
